import redis
import time
import oss2
import os
import json
import pymysql as mysql
from PIL import Image
import io
import requests
import socket
import base64


model_id = "/data/stable-diffusion/stable-diffusion-v1-4"
device = "cuda"

r = redis.StrictRedis(host='r-j6cmt5ayr2nq1m6zpp.redis.rds.aliyuncs.com', port=6379, db=0, password='r-j6cmt5ayr2nq1m6zpp:qpalmz1209@L')

conn = mysql.connect(host='rm-j6ctyd19275tjf7x1.mysql.rds.aliyuncs.com', user='root', password='qwemnb0011@', database='aipainter')
cur = conn.cursor()

ossAuth = oss2.Auth('LTAI5tMzkcDHweKw7762CzQn', '6cfufuTuoQpdwR7mckhxITwkLBWtcH')
ossBucket = oss2.Bucket(ossAuth, 'oss-cn-hongkong-internal.aliyuncs.com', 'aipainter-overseas')

#pipe = StableDiffusionPipeline.from_pretrained(model_id, revision="fp16", torch_dtype=torch.float16)
#pipe = pipe.to(device)
print('ready')
# pipeImg = StableDiffusionImg2ImgPipeline.from_pretrained(
#     model_id,
#     revision="fp16",
#     torch_dtype=torch.float16
# )
#
# pipeImg = pipeImg.to(device)
#
# print('img ready')

#hostIp = socket.gethostbyname(socket.gethostname())

while True:
    #r.hset("painter:host", hostIp, 1)
    #r.setex("painter:host:" + hostIp, 60, time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    sdinput = r.lpop("painter:queue1")
    if sdinput == None:
        cur.execute('select 1')
        time.sleep(1)
        continue
    # 获取队列数据
    print(str(sdinput,encoding="utf-8"))
    arr = json.loads(str(sdinput, encoding="utf-8"))
    taskKey = "painter:task:" + arr['task_id']
    userLenKey = "painter:user:len:" + str(arr['user_id'])
    userDotCount = "painter:opus:remaining:" + time.strftime("%y%m%d", time.localtime())

    #是否是延时任务
    delay_task = False
    if "is_dealy" in arr and arr["is_delay"] == 1:
        delay_task = True

    # 设置处理中状态
    r.hset(taskKey,'status',1)

    nsfw = False

    filename = "%s.png" % arr['task_id']
#    if arr['image'] == '':

    if arr['type'] == 1:
        # 普通sd
        payload = {
            "drawing_id": arr['drawing_id'],
            "prompt": arr['translator_prompt'],
            "steps": 20,
            "cfg_scale": "7.0",
            "width": arr['width'],
            "height": arr['height'],
            "negative_prompt": arr['translator_negative_prompt'],
            "sampler_index": "Euler a"
        }

        if arr['drawing_id'] == 6:
            #幻想画板
            headers = {"User-Agent":"illusion"}
        elif arr['drawing_id'] == 3:
            #midjourney画板
            headers = {"User-Agent":"mj"}
        elif arr['drawing_id'] == 4:
            #照片画板
            headers = {"User-Agent":"photo"}
        else:
            #经典画板
            headers = {"User-Agent":"sd15"}
                
        payload_json = json.dumps(payload)
        response = requests.post(url=f'http://region-3.seetacloud.com:33157/sdapi/v1/txt2img', data=payload_json, headers=headers).json()

        image = Image.open(io.BytesIO(base64.b64decode(response['images'][0])))
        image.save(filename)
        if response['images'][0].find('AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA') > -1 :
            nsfw = True
    if arr['type'] == 2:
        # 二次元普通生成
        payload = {
            "drawing_id": arr['drawing_id'],
            "prompt": arr['translator_prompt'],
            "steps": 20,
            "cfg_scale": "7.0",
            "width": arr['width'],
            "height": arr['height'],
            "negative_prompt": arr['translator_negative_prompt'],
            "sampler_index": "DPM++ SDE Karras"
        }
        headers = {"User-Agent":"novelai"}
        payload_json = json.dumps(payload)
        response = requests.post(url=f'http://region-3.seetacloud.com:33157/sdapi/v1/txt2img', data=payload_json, headers=headers).json()
        image = Image.open(io.BytesIO(base64.b64decode(response['images'][0])))
        image.save(filename)
        if response['images'][0].find('AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA') > -1 :
            nsfw = True
    if arr['type'] == 3:
        #先生成词
        payloadPro = {
            "image": arr['drawing_image'][0]
        }
        headers = {"User-Agent":"novelai"}
        payload_json_pro = json.dumps(payloadPro)
        responsePro = requests.post(url=f'http://region-3.seetacloud.com:33157/sdapi/v1/interrogate', data=payload_json_pro, headers=headers).json()
        if "caption" not in responsePro:
            responsePro['caption'] = "";


        # 二次元照片生成
        payload = {
            "init_images": arr['drawing_image'],
            "prompt": responsePro['caption'],
            "steps": 20,
            "cfg_scale": "7.0",
            "width": arr['width'],
            "height": arr['height'],
            "negative_prompt": arr['translator_negative_prompt'],
            "sampler_index": "DPM++ SDE Karras"
        }
        headers = {"User-Agent":"novelai"}
        payload_json = json.dumps(payload)
        response = requests.post(url=f'http://region-3.seetacloud.com:33157/sdapi/v1/img2img', data=payload_json, headers=headers).json()
        image = Image.open(io.BytesIO(base64.b64decode(response['images'][0])))
        image.save(filename)
        if response['images'][0].find('AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA') > -1 :
            nsfw = True        

    #默认值drawing_image
    if "drawing_url" not in arr:
        drawing_image = "";
    else: 
        drawing_image = arr['drawing_url'];


    # 上传图片&删除本地图片
    path = "sd/opus/%s/%s" % (time.strftime("%Y-%m-%d", time.localtime()), filename)
    ossBucket.put_object_from_file(path, filename)
    os.remove(filename)

    # 写入数据&更新状态
    url = 'https://overseas-files.dalleai.art/' + path

    if nsfw :
        url = 'https://overseas-files.dalleai.art/default.png'
        if int(arr['dot']) > 0 :
            sql = "update ap_user set dot = dot + %s where id = %s"
            cur.execute(sql, (arr['dot'], arr['user_id']))
            sql = "insert ap_dot_log(user_id, value, type) value(%s, %s, %s)"
            cur.execute(sql, (arr['user_id'], arr['dot'], 5))
            url = 'https://overseas-files.dalleai.art/default-dot.png'
            r.hincrby(userDotCount, arr['user_id'], -1)

    sql = """
    insert into ap_opus(user_id, prompt, task_id, url, height, width, title, negative_prompt, type, use_dot, drawing_id, drawing_image) value(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
    """
    cur.execute(sql, (arr['user_id'], arr['prompt'], arr['task_id'], url, arr['height'], arr['width'], arr['prompt'], arr['translator_negative_prompt'], arr['type'],arr['dot'],arr['drawing_id'],drawing_image))
    conn.commit()
    if not delay_task:
        r.hset(taskKey,'status',2)
        r.hset(taskKey,'url',url)
        r.hset(taskKey,'prompt',arr['prompt'])
    r.decrby(userLenKey, 1)
    r.expire(userLenKey, 600)

    # 打印完成
    print(url)

